[1/2] refactor(speculative): hand the DFlash-family loss one context per training step - #2989

…nt and head

HFDominoModel did not override pseudo_speculative_generate, so AR
validation ran DFlash's, which decodes block positions 1..bs-1 with the
backbone alone. DFlash trains position k on the token at anchor + k;
Domino trains it on anchor + k + 1 (_label_shift = 1). Every drafted
token was therefore compared with the target one position too early,
so drafts were almost never accepted, and the correction head that
training optimizes and the exporter ships was never applied. domino.yaml
keeps AR validation off, so shipped training did not hit this.

DFlash's default decode now starts at block position 1 - _label_shift
and returns at most bs - 1 + _label_shift tokens. For DFlash, DFlash2
and LiLiCorr (shift 0) that is the same slice as before. Domino
overrides it the way DSpark does for its Markov head: it decodes the
block left to right, stepping the GRU over the anchor and the tokens
drafted so far, and from pure_draft_prefix_len on adds
embed_proj([hidden_k; GRU state]) to the backbone logits, as
_apply_domino_head does in training. As in DSpark's decode, the
backbone logits and the drafted tokens' embeddings are moved to the
draft's device: under device_map="auto" the embedding table, and a tied
lm_head with it, can sit on another GPU.

The eval forward warned that Domino's acceptance rates were
backbone-only, but that method computes no acceptance rate, and AR
validation now runs the head. The warning goes, and so do the other
copies of the claim: the comments in domino.yaml, dspark.yaml and the
Nemotron-3.5 DSpark warm-start recipe, the note in the Domino launcher
example, and the speculative-decoding skill's Domino, DSpark and
validation sheets. DSpark's generation has applied its Markov head
since DSpark was added. The validation sheet keeps its rule against
reporting such an AR as the model's result, now for the heads that AR
validation still skips (LiLiCorr's reranker, DFlash2's selector), and
adds the LiLiCorr and online DFlash2 examples to the ones it lists as
shipping no benchmark task. The Domino sheet and example now say how to
add the HF AR eval: DFlash's task_2 with --steps equal to the block
size, since Domino drafts the whole block. The flag values do not
change.

Verified: a new test checks that Domino's greedy drafts equal the
argmax of _apply_domino_head teacher-forced on them, with head weights
large enough to overturn the backbone. It fails on the parent commit,
and also when the head is skipped, its GRU state lags a step, or the
prefix boundary is off by one. The GPU test that puts lm_head on
another device now covers Domino as well, and fails without the move.
Bitwise A/B against #2757: DFlash and DSpark loss, gradients and drafted
tokens are unchanged; only Domino's drafted tokens change.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>

DSpark and Domino train block position k on the token at anchor + k + 1
(HFDFlashModel._label_shift = 1). The same fact was also written as
DSparkModule.shift_label = True, as DominoModule.shift_label read from
the config, as a NotImplementedError in HFDominoModel.modify, and as the
exporters' getattr(draft_config, "shift_label", True), which read the
user's raw architecture config.

A DSpark config with shift_label: false therefore trained with the shift
but exported shift_label: false, so a loader that honors the field would
read every drafted position one token off. Domino refused the same
config, but only after building the draft.

HFDFlashModel.modify now refuses shift_label: false for every variant
whose _label_shift is 1, before anything is built; unset or true
converts as before. Both exporters write shift_label from _label_shift,
and the vanilla DSpark export copies it into sample_from_anchor, vLLM's
name for the same alignment, instead of writing a second literal True.
The modules' shift_label attributes and Domino's own check go, since
nothing reads them any more. Domino's error becomes a ValueError, like
the other convert-time config checks, and the skill's Domino and DSpark
sheets list it.

The refusal applies at convert only. Restore replays convert, so
restore_dflash_model drops shift_label from the saved architecture
config, as it already drops the warm-start path. A DSpark checkpoint
saved with shift_label: false trained next-token alignment all the same:
it still loads, and now re-exports shift_label: true.

Verified: a new test converts DSpark and Domino with shift_label unset,
as the recipes leave it, restores a modelopt_state that carries
shift_label: false, and expects shift_label: false to be refused at
convert. On the parent commit DSpark does not raise and Domino raises
NotImplementedError; without the restore change the restore raises.
A tiny DSpark checkpoint saved by 0.47.0 with shift_label: false loads
through from_pretrained and exports shift_label: true. The existing
export tests still read shift_label: true and, for vanilla heads,
sample_from_anchor: true.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>

… was built with

DSparkModule resolves its head fields from the architecture config as
it builds them: it lower-cases markov_head_type, casts markov_rank to
int and use_confidence_head to bool, and fills in defaults for the ones
left out. #2902 made DSparkExporter read markov_head_type off the built
module, but markov_rank and use_confidence_head were still read from
the raw config dict, the latter with a second default. A markov_rank
given as a float (64.0) therefore trained a rank-64 head but exported
64.0, and the vanilla export copies that value to the top level, where
vLLM's Qwen3DSparkModel sizes its Markov table from it. DominoExporter
likewise re-derived pure_draft_prefix_len with a second default.

Both exporters now read every head field off the built module, as
LiLiCorrExporter already does for its reranker, so the exported config
describes the head that the saved weights belong to.

Verified: the DSpark export-config test now converts with markov_rank
given as a float and expects an int at both levels of config.json; all
eight cases fail on the parent commit. Domino's exported values and
DSpark's use_confidence_head are unchanged for every config, since the
modules store them as given or cast them the same way.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>

DSpark and Domino have shipped since 0.46.0, so these are fixes to
released features, and two of them ask something of users: drop
shift_label: false from DSpark configs, which now fail to convert, and
re-export DSpark drafters whose export carried shift_label: false. The
entry leaves out markov_head_type: #2902 already exports it as built.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>

…step

HFDFlashModel.forward built the batch, ran the draft and computed the loss inline, and the
variants built on it got at the loss's pieces piecemeal: DFlash2 asked _compute_loss for
its unreduced terms (return_terms, _DFlashLossTerms) to recompose the block objective, and
DFlash2 and LiLiCorr overrode forward() only to carry out the metrics their _compute_loss
had left on the model (_selector_metrics, _lilicorr_metrics).

The forward now hands the loss one object per step:
- DraftBatch: the batch, its blocks and the draft's hidden states.
- LossContext: the labels (_block_targets), the draft's logits (DraftLogits: the
  distribution the draft serves, and the backbone's before any head) and the position
  weights, with the per-token cross-entropy and the teacher's logits computed lazily, once.
- _compute_loss(ctx) returns a LossOutput: the loss, the accuracy train_acc reports and
  detached 0-dim metrics, which the forward returns as draft_metrics.

What a variant may change becomes a hook: _online_base_outputs (the base model's outputs),
_draft_logits (a head on the backbone's logits), _position_weighting and _compute_loss.
The loss mask comes from _build_loss_mask. The DSpark/Domino padding fix, a separate PR
against main, adds the same helper, with the same text, at the same place (right after
get_exporter), so a merge of the two keeps one copy of it.

DFlash2 and LiLiCorr take the context: DFlash2's tv/lambda blend reads the cross-entropy
and weights from it, and both return their metrics instead of overriding forward(). DSpark
and Domino keep their own forwards for now.

Verified against the previous commit with one train step and one generation step per case,
in fp32 and bf16: DFlash (KD and CE), DFlash2 (ce/tv/lambda with the selector) and LiLiCorr,
offline and online, D-PACE and decay, including zero-loss steps, give bitwise-equal loss,
gradients, train_acc and generated tokens (32/32). The selector and LiLiCorr metrics keep
their names and values under draft_metrics.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>

…sses.py

The block-supervision helpers (_block_positions, _block_gather, _block_tokens,
_BlockTargets, _block_targets), the position weighting (_dpace_position_weights,
_block_position_weights), _masked_top1, the loss context (DraftLogits, LossOutput,
LossContext) and DSpark's per-token TVD (_tvd_chunk, _tvd_per_token) move unchanged from
hf_dflash.py and hf_dspark.py into one module that names no variant. The next PR in this
stack builds the shared loss terms on them, and hf_dspark.py would otherwise be imported by
the generic loss code for its TVD.

A pure move: every moved definition is byte-identical to its previous text, and only
import lines change elsewhere. The new file carries the SpecForge and DeepSpec notices of
the two files the code came from and is excluded from the license-insertion hook like them.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>

…_losses.py

Update the call sites of the internals the previous two commits changed or moved, keeping
each test's intent:
- _compute_loss tests (D-PACE/decay/KD branches, trailing padding blocks, KD from the base
  hidden vs. full logits) go through _loss_context and _compute_loss on given draft logits.
- "KD without base_outputs raises" becomes "KD without a teacher raises": the context
  always carries the base outputs, so the remaining failure is a batch with neither the
  base logits nor its final hidden.
- The selector and LiLiCorr metrics are read from draft_metrics, LiLiCorr's as tensors.
- _block_targets, _dpace_position_weights and DSpark's TVD helpers are imported from
  dflash_losses, where the move put them.

One new test pins which objectives keep the per-token cross-entropy's gradient. The
context computes the CE once and sets its grad mode from ce_grad, overriding the no_grad
under which D-PACE reads it first. The test inspects the context that a real DFlash2
training forward under D-PACE hands _compute_loss: ce, tv and lambda keep the CE's
gradient, and KD does not.

Mutation check. If _loss_context always asks for the gradient, LossContext ignores
ce_grad, or the CE always keeps its graph, then each KD step under D-PACE (the default
configuration) holds a vocab-wide graph that it never backpropagates. The loss and
gradients stay bitwise equal (DFlash, DFlash2 and LiLiCorr KD, fp32 and bf16: 6/6), and
the rest of the suite passes (512 passed); only [kd] fails. Dropping the gradient instead
already failed three existing tests, and now also fails [ce], [tv] and [lambda].

tests/unit/torch/speculative and test_hf_spec_rope_export.py: 513 passed on transformers
5.5.4, 5.12.1 and 5.18.0. On 4 CPUs, the changed tests take at most 3.4 s each and the new
cases under 0.1 s.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>

…ing changes

DFlash2's selector_metrics and LiLiCorr's lilicorr_metrics arrive in 0.48.0, which this
change follows. The first commit of this series replaces both keys with draft_metrics, keeps
their inner keys, and returns tensors where LiLiCorr returned Python floats. Nothing in the
repo reads these keys, but user code that reads them breaks.

The entry goes under the 0.49.0 heading: release/0.48.0 is branched and ships both keys,
so this change first reaches users in 0.49.0.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>